// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <bitset>
#include <pollux/common/base/gtest_macros.h>
#include <pollux/dwio/common/data_buffer.h>

namespace kumo::pollux::dwio::common {
    namespace {
        constexpr uint32_t DEFAULT_MAX_PAGE_BYTES = 256 * 1024;
        constexpr uint32_t DEFAULT_INITIAL_CAPACITY = 1024;
    } // namespace

    template<typename T>
    class ChainedBuffer {
    public:
        explicit ChainedBuffer(
            pollux::memory::MemoryPool &pool,
            uint32_t initialCapacity = DEFAULT_INITIAL_CAPACITY,
            uint32_t maxPageBytes = DEFAULT_MAX_PAGE_BYTES)
            : pool_{pool},
              size_{0},
              capacity_{0},
              initialCapacity_{initialCapacity},
              maxPageCapacity_{maxPageBytes / static_cast<uint32_t>(sizeof(T))},
              bits_{trailingZeros(maxPageCapacity_)},
              mask_{~(~0ull << bits_)} {
            POLLUX_CHECK_EQ(
                bitCount(maxPageBytes), 1, "must be power of 2: {}", maxPageBytes);
            POLLUX_CHECK_EQ(bitCount(sizeof(T)), 1, "must be power of 2: {}", sizeof(T));

            while (capacity_ < initialCapacity_) {
                addPage(std::min(initialCapacity_, maxPageCapacity_));
            }
        }

        T &operator[](uint64_t index) {
            return getPageUnsafe(index)[getPageOffset(index)];
        }

        const T &operator[](uint64_t index) const {
            return getPageUnsafe(index)[getPageOffset(index)];
        }

        void reserve(uint32_t capacity) {
            if (MELON_LIKELY(capacity <= capacity_)) {
                return;
            }
            // capacity grow rules
            // 1. if capacity < pageCapacity (only first page is needed), grow to next
            // power of 2
            // 2. if capacity > pageCapacity, allocate pages until it fits
            const auto pageCount = ((capacity - 1) / maxPageCapacity_) + 1;

            if (pages_.empty()) {
                addPage(pageCount == 1 ? nextPowOf2(capacity) : maxPageCapacity_);
            } else if (pages_[0].capacity() < maxPageCapacity_) {
                resizeFirstPage(pageCount == 1 ? nextPowOf2(capacity) : maxPageCapacity_);
            }

            while (pages_.size() < pageCount) {
                addPage(maxPageCapacity_);
            }
        }

        void append(T t) {
            reserve(size_ + 1);
            unsafeAppend(t);
        }

        void unsafeAppend(T t) {
            (*this)[size_++] = t;
        }

        uint64_t capacity() const {
            return capacity_;
        }

        uint64_t size() const {
            return size_;
        }

        /// If 'all' is true, clear all pages. Otherwise, clear all the pages except
        /// the first one.
        void clear(bool all = false) {
            if (pages_.empty()) {
                POLLUX_CHECK_EQ(capacity_, 0);
                POLLUX_CHECK_EQ(size_, 0);
                return;
            }

            if (all) {
                pages_.clear();
                capacity_ = 0;
            } else {
                while (pages_.size() > 1) {
                    pages_.pop_back();
                }
                capacity_ = pages_.back().capacity();
            }
            size_ = 0;
        }

        void applyRange(
            uint64_t begin,
            uint64_t end,
            std::function<void(const T *, uint64_t, uint64_t)> fn) {
            POLLUX_CHECK_LE(begin, end);
            POLLUX_CHECK_LE(end, size_);
            auto pageBegin = getPageIndex(begin);
            auto indexBegin = getPageOffset(begin);
            auto pageEnd = getPageIndex(end);
            auto indexEnd = getPageOffset(end);
            while (pageBegin < pageEnd ||
                   (pageBegin == pageEnd && indexBegin < indexEnd)) {
                fn(pages_.at(pageBegin).data(),
                   indexBegin,
                   pageBegin == pageEnd ? indexEnd : maxPageCapacity_);
                ++pageBegin;
                indexBegin = 0;
            }
        }

    private:
        void addPage(uint32_t capacity) {
            POLLUX_CHECK_LE(capacity, maxPageCapacity_);
            // Either there is no page or the first page is at its max capacity
            POLLUX_CHECK(
                pages_.empty() ||
                (pages_[0].capacity() == maxPageCapacity_ &&
                    capacity == maxPageCapacity_));

            pages_.emplace_back(pool_);
            pages_.back().reserve(capacity);
            capacity_ += capacity;
        }

        void resizeFirstPage(uint32_t capacity) {
            POLLUX_CHECK(
                pages_.size() == 1 && capacity <= maxPageCapacity_ &&
                capacity_ < capacity);

            pages_.back().reserve(capacity);
            capacity_ = capacity;
        }

        DataBuffer<T> &getPageUnsafe(uint64_t index) {
            return pages_[getPageIndex(index)];
        }

        const DataBuffer<T> &getPageUnsafe(uint64_t index) const {
            return pages_[getPageIndex(index)];
        }

        uint64_t getPageIndex(uint64_t index) const {
            return index >> bits_;
        }

        uint64_t getPageOffset(uint64_t index) const {
            return index & mask_;
        }

        static size_t bitCount(uint32_t val) {
            return std::bitset<64>(val).count();
        }

        static uint8_t trailingZeros(uint32_t val) {
            POLLUX_CHECK_GT(val, 0);
            uint8_t ret = 0;
            val = (val ^ (val - 1)) >> 1;
            while (val) {
                val >>= 1;
                ++ret;
            }
            return ret;
        }

        static uint32_t nextPowOf2(uint32_t val) {
            val -= 1;
            val |= (val >> 1);
            val |= (val >> 2);
            val |= (val >> 4);
            val |= (val >> 8);
            val |= (val >> 16);
            return val + 1;
        }

        pollux::memory::MemoryPool &pool_;
        std::vector<DataBuffer<T> > pages_;
        uint64_t size_;
        uint64_t capacity_;

        const uint32_t initialCapacity_;
        const uint32_t maxPageCapacity_;
        const uint8_t bits_;
        const uint64_t mask_;

        POLLUX_FRIEND_TEST(ChainedBufferTests, testCreate);
        POLLUX_FRIEND_TEST(ChainedBufferTests, testReserve);
        POLLUX_FRIEND_TEST(ChainedBufferTests, testAppend);
        POLLUX_FRIEND_TEST(ChainedBufferTests, testClear);
        POLLUX_FRIEND_TEST(ChainedBufferTests, testGetPage);
        POLLUX_FRIEND_TEST(ChainedBufferTests, testGetPageIndex);
        POLLUX_FRIEND_TEST(ChainedBufferTests, testGetPageOffset);
        POLLUX_FRIEND_TEST(ChainedBufferTests, testBitCount);
        POLLUX_FRIEND_TEST(ChainedBufferTests, testTrailingZeros);
        POLLUX_FRIEND_TEST(ChainedBufferTests, testClearAll);
    };
} // namespace kumo::pollux::dwio::common
